use burn_import::onnx::{ModelGen, RecordType};

fn main() {
    // Re-run this build script if the onnx-tests directory changes.
    println!("cargo:rerun-if-changed=tests");

    // Add onnx models.
    ModelGen::new()
        .input("tests/abs/abs.onnx")
        .input("tests/add/add.onnx")
        .input("tests/add/add_shape.onnx")
        .input("tests/add/add_broadcast.onnx")
        .input("tests/initializer_to_const/initializer_to_const.onnx")
        .input("tests/and/and.onnx")
        .input("tests/and/and_broadcast.onnx")
        .input("tests/and/and_scalar.onnx")
        .input("tests/add/add_int.onnx")
        .input("tests/add/add_shape_tensor.onnx")
        .input("tests/add/add_argmax_with_shape.onnx")
        .input("tests/argmax/argmax.onnx")
        .input("tests/argmax/argmax_both_keepdims.onnx")
        .input("tests/argmax/argmax_1d.onnx")
        .input("tests/argmin/argmin.onnx")
        .input("tests/argmin/argmin_both_keepdims.onnx")
        .input("tests/argmin/argmin_1d.onnx")
        .input("tests/attention/attention_4d.onnx")
        .input("tests/attention/attention_3d.onnx")
        .input("tests/attention/attention_attn_mask_bool.onnx")
        .input("tests/attention/attention_attn_mask_int.onnx")
        .input("tests/attention/attention_attn_mask_float.onnx")
        .input("tests/attention/attention_softcap.onnx")
        .input("tests/attention/attention_cache.onnx")
        .input("tests/attention/attention_custom_scale.onnx")
        .input("tests/attention/attention_is_causal.onnx")
        .input("tests/attention/attention_qk_output_0.onnx")
        .input("tests/attention/attention_qk_output_1.onnx")
        .input("tests/attention/attention_qk_output_2.onnx")
        .input("tests/attention/attention_qk_output_3.onnx")
        .input("tests/avg_pool1d/avg_pool1d.onnx")
        .input("tests/avg_pool2d/avg_pool2d.onnx")
        .input("tests/batch_norm/batch_norm.onnx")
        .input("tests/bitshift/bitshift_left.onnx")
        .input("tests/bitshift/bitshift_left_scalar.onnx")
        .input("tests/bitshift/scalar_bitshift_left.onnx")
        .input("tests/bitshift/scalar_bitshift_left_scalar.onnx")
        .input("tests/bitshift/bitshift_right.onnx")
        .input("tests/bitshift/bitshift_right_scalar.onnx")
        .input("tests/bitshift/scalar_bitshift_right.onnx")
        .input("tests/bitshift/scalar_bitshift_right_scalar.onnx")
        .input("tests/bitwise_and/bitwise_and.onnx")
        .input("tests/bitwise_and/bitwise_and_scalar.onnx")
        .input("tests/bitwise_and/scalar_bitwise_and.onnx")
        .input("tests/bitwise_and/scalar_bitwise_and_scalar.onnx")
        .input("tests/bitwise_not/bitwise_not.onnx")
        .input("tests/bitwise_or/bitwise_or.onnx")
        .input("tests/bitwise_or/bitwise_or_scalar.onnx")
        .input("tests/bitwise_or/scalar_bitwise_or.onnx")
        .input("tests/bitwise_or/scalar_bitwise_or_scalar.onnx")
        .input("tests/bitwise_xor/bitwise_xor.onnx")
        .input("tests/bitwise_xor/bitwise_xor_scalar.onnx")
        .input("tests/bitwise_xor/scalar_bitwise_xor.onnx")
        .input("tests/bitwise_xor/scalar_bitwise_xor_scalar.onnx")
        .input("tests/bernoulli/bernoulli.onnx")
        .input("tests/cast/cast.onnx")
        .input("tests/cast/cast_shape.onnx")
        .input("tests/cast/cast_shape_to_float.onnx")
        .input("tests/cast/cast_shape_to_bool.onnx")
        .input("tests/ceil/ceil.onnx")
        .input("tests/clip/clip.onnx")
        .input("tests/concat/concat.onnx")
        .input("tests/concat/concat_shape.onnx")
        .input("tests/concat/concat_shape_with_constant.onnx")
        .input("tests/concat/concat_mixed_single_element.onnx")
        .input("tests/concat/concat_mixed_three_elements.onnx")
        .input("tests/concat/concat_multiple_mixed.onnx")
        .input("tests/concat/concat_with_constants.onnx")
        .input("tests/constant/constant_f32.onnx")
        .input("tests/constant/constant_f64.onnx")
        .input("tests/constant/constant_i32.onnx")
        .input("tests/constant/constant_i64.onnx")
        .input("tests/constant/constant_bool.onnx")
        .input("tests/constant/constant_shape.onnx")
        .input("tests/constant/constant_tensor_f32.onnx")
        .input("tests/constant/constant_tensor_i32.onnx")
        .input("tests/constant/constant_tensor_bool.onnx")
        .input("tests/constant/constant_empty_tensor_f32.onnx")
        .input("tests/constant/rank_inference_propagation.onnx")
        .input("tests/constant/shape_binary_ops_with_constant.onnx")
        .input("tests/constant_of_shape/constant_of_shape.onnx")
        .input("tests/constant_of_shape/constant_of_shape_full_like.onnx")
        .input("tests/constant_of_shape/constant_of_shape_scalar.onnx")
        .input("tests/constant_of_shape/constant_of_shape_scalar_custom_value.onnx")
        .input("tests/constant_of_shape/constant_of_shape_tensor.onnx")
        .input("tests/constant_of_shape/constant_of_shape_shape_optimization.onnx")
        .input("tests/constant_of_shape/constant_of_shape_with_constant_input.onnx")
        .input("tests/constant_lifting_multiple/constant_lifting_multiple.onnx")
        .input("tests/constant_lifting_multiple/constant_reused.onnx")
        .input("tests/conv1d/conv1d.onnx")
        .input("tests/conv2d/conv2d.onnx")
        .input("tests/conv3d/conv3d.onnx")
        .input("tests/conv_transpose1d/conv_transpose1d.onnx")
        .input("tests/conv_transpose2d/conv_transpose2d.onnx")
        .input("tests/conv_transpose3d/conv_transpose3d.onnx")
        .input("tests/cos/cos.onnx")
        .input("tests/cosh/cosh.onnx")
        .input("tests/depth_to_space/depth_to_space_dcr.onnx")
        .input("tests/depth_to_space/depth_to_space_crd.onnx")
        .input("tests/div/div.onnx")
        .input("tests/div/div_shape.onnx")
        .input("tests/div/div_shape_tensor.onnx")
        .input("tests/div/div_broadcast.onnx")
        .input("tests/mod/modulo.onnx")
        .input("tests/mod/mod_scalar.onnx")
        .input("tests/mod/mod_remainder.onnx")
        .input("tests/mod/mod_fmod.onnx")
        .input("tests/mod/mod_broadcast_fixed.onnx")
        .input("tests/mod/mod_broadcast_remainder_fixed.onnx")
        .input("tests/dropout/dropout.onnx")
        .input("tests/empty_graph/empty_graph_scalar.onnx")
        .input("tests/empty_graph/empty_graph_scalar_int.onnx")
        .input("tests/empty_graph/empty_graph_shape.onnx")
        .input("tests/empty_graph/empty_graph_tensor.onnx")
        .input("tests/empty_graph/empty_graph_multiple.onnx")
        .input("tests/equal/equal.onnx")
        .input("tests/equal/equal_shape.onnx")
        .input("tests/equal/equal_two_shapes.onnx")
        .input("tests/erf/erf.onnx")
        .input("tests/exp/exp.onnx")
        .input("tests/expand/expand.onnx")
        .input("tests/expand/expand_tensor.onnx")
        .input("tests/expand/expand_shape.onnx")
        .input("tests/expand/expand_with_where_shape.onnx")
        .input("tests/eye_like/eye_like.onnx")
        .input("tests/eye_like/eye_like_k1.onnx")
        .input("tests/eye_like/eye_like_int.onnx")
        .input("tests/eye_like/eye_like_k_minus1.onnx")
        .input("tests/eye_like/eye_like_float64.onnx")
        .input("tests/eye_like/eye_like_int32.onnx")
        .input("tests/eye_like/eye_like_bool.onnx")
        .input("tests/eye_like/eye_like_large_k.onnx")
        .input("tests/eye_like/eye_like_1x1.onnx")
        .input("tests/eye_like/eye_like_wide.onnx")
        .input("tests/eye_like/eye_like_neg_large_k.onnx")
        .input("tests/nonzero/nonzero_float32.onnx")
        .input("tests/nonzero/nonzero_int64.onnx")
        .input("tests/nonzero/nonzero_bool.onnx")
        .input("tests/nonzero/nonzero_1d.onnx")
        .input("tests/nonzero/nonzero_3d.onnx")
        .input("tests/nonzero/nonzero_empty.onnx")
        .input("tests/flatten/flatten.onnx")
        .input("tests/flatten/flatten_2d.onnx")
        .input("tests/floor/floor.onnx")
        .input("tests/gather/gather_1d_idx.onnx")
        .input("tests/gather/gather_2d_idx.onnx")
        .input("tests/gather/gather_scalar.onnx")
        .input("tests/gather/gather_constant_2d_indices.onnx")
        .input("tests/gather/gather_static_shape_indices.onnx")
        .input("tests/gather/gather_shape.onnx")
        .input("tests/gather/gather_with_shape_indices.onnx")
        .input("tests/gather/gather_scalar_out.onnx")
        .input("tests/gather_elements/gather_elements.onnx")
        .input("tests/gelu/gelu.onnx")
        .input("tests/gemm/gemm.onnx")
        .input("tests/gemm/gemm_non_unit_alpha_beta.onnx")
        .input("tests/gemm/gemm_no_c.onnx")
        .input("tests/global_avr_pool/global_avr_pool.onnx")
        .input("tests/graph_multiple_output_tracking/graph_multiple_output_tracking.onnx")
        .input("tests/greater/greater.onnx")
        .input("tests/greater/greater_scalar.onnx")
        .input("tests/greater/greater_broadcast.onnx")
        .input("tests/greater_or_equal/greater_or_equal.onnx")
        .input("tests/greater_or_equal/greater_or_equal_scalar.onnx")
        .input("tests/greater_or_equal/greater_or_equal_broadcast.onnx")
        .input("tests/group_norm/group_norm.onnx")
        .input("tests/hard_sigmoid/hard_sigmoid.onnx")
        .input("tests/identity/identity_constant.onnx")
        .input("tests/identity/identity_passthrough.onnx")
        .input("tests/identity/identity_chain.onnx")
        .input("tests/identity/identity_only.onnx")
        .input("tests/instance_norm1d/instance_norm1d.onnx")
        .input("tests/instance_norm2d/instance_norm2d.onnx")
        .input("tests/instance_norm3d/instance_norm3d.onnx")
        .input("tests/is_inf/is_inf.onnx")
        .input("tests/is_inf/is_inf_scalar.onnx")
        .input("tests/is_inf/is_inf_neg_only.onnx")
        .input("tests/is_inf/is_inf_pos_only.onnx")
        .input("tests/is_inf/is_inf_none.onnx")
        .input("tests/is_nan/is_nan.onnx")
        .input("tests/is_nan/is_nan_scalar.onnx")
        .input("tests/layer_norm/layer_norm.onnx")
        .input("tests/leaky_relu/leaky_relu.onnx")
        .input("tests/less/less.onnx")
        .input("tests/less/less_scalar.onnx")
        .input("tests/less/less_broadcast.onnx")
        .input("tests/less_or_equal/less_or_equal.onnx")
        .input("tests/less_or_equal/less_or_equal_scalar.onnx")
        .input("tests/less_or_equal/less_or_equal_broadcast.onnx")
        .input("tests/linear/linear.onnx")
        .input("tests/log/log.onnx")
        .input("tests/log_softmax/log_softmax.onnx")
        .input("tests/where_op/where_op.onnx")
        .input("tests/where_op/where_op_broadcast.onnx")
        .input("tests/where_op/where_op_scalar_x.onnx")
        .input("tests/where_op/where_op_scalar_y.onnx")
        .input("tests/where_op/where_op_all_scalar.onnx")
        .input("tests/where_op/where_shape_all_shapes.onnx")
        .input("tests/where_op/where_shape_scalar_cond.onnx")
        .input("tests/where_op/where_shapes_from_inputs.onnx")
        .input("tests/where_op/where_static_shape.onnx")
        .input("tests/matmul/matmul.onnx")
        .input("tests/matmulinteger/matmulinteger.onnx")
        .input("tests/matmulinteger/matmulinteger_ranks.onnx")
        .input("tests/matmul/matmul_ranks.onnx")
        .input("tests/max/max.onnx")
        .input("tests/maxpool1d/maxpool1d.onnx")
        .input("tests/maxpool2d/maxpool2d.onnx")
        .input("tests/min/min.onnx")
        .input("tests/mean/mean.onnx")
        .input("tests/mul/mul.onnx")
        .input("tests/mul/mul_shape.onnx")
        .input("tests/mul/mul_shape_tensor.onnx")
        .input("tests/mul/mul_broadcast.onnx")
        .input("tests/neg/neg.onnx")
        .input("tests/not/not.onnx")
        .input("tests/one_hot/one_hot.onnx")
        .input("tests/or/or.onnx")
        .input("tests/or/or_scalar.onnx")
        .input("tests/or/or_broadcast.onnx")
        .input("tests/pad/pad.onnx")
        .input("tests/pow/pow.onnx")
        .input("tests/pow/pow_int.onnx")
        .input("tests/prelu/prelu.onnx")
        .input("tests/random_normal/random_normal.onnx")
        .input("tests/random_normal_like/random_normal_like.onnx")
        .input("tests/random_uniform/random_uniform.onnx")
        .input("tests/random_uniform_like/random_uniform_like.onnx")
        .input("tests/range/range.onnx")
        .input("tests/range/range_static.onnx")
        .input("tests/range/range_mixed.onnx")
        .input("tests/range/range_runtime.onnx")
        .input("tests/recip/recip.onnx")
        .input("tests/reduce/reduce_max.onnx")
        .input("tests/reduce/reduce_max_bool.onnx")
        .input("tests/reduce/reduce_mean.onnx")
        .input("tests/reduce/reduce_mean_partial_shape.onnx")
        .input("tests/reduce/reduce_min.onnx")
        .input("tests/reduce/reduce_min_bool.onnx")
        .input("tests/reduce/reduce_prod.onnx")
        .input("tests/reduce/reduce_sum.onnx")
        .input("tests/reduce/reduce_sum_square.onnx")
        .input("tests/reduce/reduce_l1.onnx")
        .input("tests/reduce/reduce_l2.onnx")
        .input("tests/reduce/reduce_log_sum.onnx")
        .input("tests/reduce/reduce_log_sum_exp.onnx")
        .input("tests/relu/relu.onnx")
        .input("tests/reshape/reshape.onnx")
        .input("tests/reshape/reshape_with_1d_tensor.onnx")
        .input("tests/reshape/reshape_with_shape.onnx")
        .input("tests/reshape/reshape_to_scalar.onnx")
        .input("tests/reshape/reshape_3d_to_scalar.onnx")
        .input("tests/reshape/reshape_shape_to_shape.onnx")
        .input("tests/reshape/reshape_shape_with_neg.onnx")
        .input("tests/reshape/reshape_shape_partial.onnx")
        .input("tests/resize/resize_with_sizes.onnx")
        .input("tests/resize/resize_1d_linear_scale.onnx")
        .input("tests/resize/resize_1d_nearest_scale.onnx")
        .input("tests/resize/resize_2d_bicubic_scale.onnx")
        .input("tests/resize/resize_2d_bilinear_scale.onnx")
        .input("tests/resize/resize_2d_nearest_scale.onnx")
        .input("tests/resize/resize_with_shape.onnx")
        .input("tests/resize/resize_with_sizes_tensor.onnx")
        .input("tests/round/round.onnx")
        .input("tests/shape/shape.onnx")
        .input("tests/shape/shape_of_shape.onnx")
        .input("tests/shape/shape_slice.onnx")
        .input("tests/shape/shape_chain.onnx")
        .input("tests/sigmoid/sigmoid.onnx")
        .input("tests/sign/sign.onnx")
        .input("tests/sin/sin.onnx")
        .input("tests/sinh/sinh.onnx")
        .input("tests/size/size.onnx")
        .input("tests/slice/slice.onnx")
        .input("tests/slice/slice_shape.onnx")
        .input("tests/slice/slice_scalar.onnx")
        .input("tests/slice/slice_mixed.onnx")
        .input("tests/slice/slice_shape_gather.onnx")
        .input("tests/slice/slice_shape_runtime.onnx")
        .input("tests/slice/slice_shape_multi.onnx")
        .input("tests/slice/slice_shape_negative.onnx")
        .input("tests/slice/slice_shape_negative_range.onnx")
        .input("tests/slice/slice_1d_tensor.onnx")
        .input("tests/slice/slice_shape_start_tensor_end.onnx")
        .input("tests/slice/slice_tensor_start_shape_end.onnx")
        .input("tests/slice/slice_axes.onnx")
        .input("tests/slice/slice_with_steps.onnx")
        .input("tests/slice/slice_shape_with_steps.onnx")
        .input("tests/softmax/softmax.onnx")
        .input("tests/space_to_depth/space_to_depth.onnx")
        .input("tests/sqrt/sqrt.onnx")
        .input("tests/squeeze/squeeze_multiple.onnx")
        .input("tests/squeeze/squeeze.onnx")
        .input("tests/squeeze/squeeze_shape.onnx")
        .input("tests/squeeze/squeeze_shape_noop.onnx")
        .input("tests/squeeze/squeeze_scalar.onnx")
        .input("tests/squeeze/squeeze_float.onnx")
        .input("tests/squeeze/squeeze_tensor_to_scalar.onnx")
        .input("tests/squeeze/squeeze_opset13_axes_input.onnx")
        .input("tests/squeeze/squeeze_no_axes.onnx")
        .input("tests/sub/sub.onnx")
        .input("tests/sub/sub_shape.onnx")
        .input("tests/sub/sub_broadcast.onnx")
        .input("tests/sub/sub_int.onnx")
        .input("tests/sub/sub_shape_tensor.onnx")
        .input("tests/sum/sum.onnx")
        .input("tests/sum/sum_int.onnx")
        .input("tests/tan/tan.onnx")
        .input("tests/tanh/tanh.onnx")
        .input("tests/tile/tile.onnx")
        .input("tests/topk/topk.onnx")
        .input("tests/trilu/trilu_upper.onnx")
        .input("tests/trilu/trilu_lower.onnx")
        .input("tests/transpose/transpose.onnx")
        .input("tests/unsqueeze/unsqueeze_runtime_axes.onnx")
        .input("tests/unsqueeze/unsqueeze_like.onnx")
        .input("tests/unsqueeze/unsqueeze_int_to_shape.onnx")
        .input("tests/unsqueeze/squeeze_unsqueeze_roundtrip.onnx")
        .input("tests/split/split.onnx")
        .input("tests/xor/xor.onnx")
        .input("tests/xor/xor_scalar.onnx")
        .input("tests/xor/xor_broadcast.onnx")
        // If operator tests
        .input("tests/if_op/if_conv2d.onnx")
        .input("tests/if_op/if_linear.onnx")
        .input("tests/if_op/nested_if.onnx")
        // Loop operator tests
        .input("tests/loop/loop_simple.onnx")
        .input("tests/loop/loop_dynamic_cond.onnx")
        .input("tests/loop/loop_multi_deps.onnx")
        .input("tests/loop/loop_nested.onnx")
        .input("tests/loop/loop_scan_outputs.onnx")
        // Scan operator tests
        .input("tests/scan/scan_cumsum.onnx")
        .input("tests/scan/scan_reverse.onnx")
        .input("tests/scan/scan_multi_state.onnx")
        .input("tests/scan/scan_axis1.onnx")
        // Deeply nested subgraph tests (If -> Loop -> If -> Scan)
        .input("tests/subgraph/nested_if_loop_if.onnx")
        .input("tests/subgraph/nested_if_loop_if_scan.onnx")
        .out_dir("model/")
        .run_from_script();

    // The following tests are used to generate the model with different record types.
    // (e.g. bincode, pretty_json, etc.) Do not need to add new tests here, just use the default
    // record type to the ModelGen::new() call above.

    ModelGen::new()
        .input("tests/conv1d/conv1d.onnx")
        .out_dir("model/named_mpk/")
        .record_type(RecordType::NamedMpk)
        .run_from_script();

    ModelGen::new()
        .input("tests/conv1d/conv1d.onnx")
        .out_dir("model/named_mpk_half/")
        .record_type(RecordType::NamedMpk)
        .half_precision(true)
        .run_from_script();

    ModelGen::new()
        .input("tests/conv1d/conv1d.onnx")
        .out_dir("model/pretty_json/")
        .record_type(RecordType::PrettyJson)
        .run_from_script();

    ModelGen::new()
        .input("tests/conv1d/conv1d.onnx")
        .out_dir("model/pretty_json_half/")
        .record_type(RecordType::PrettyJson)
        .half_precision(true)
        .run_from_script();

    ModelGen::new()
        .input("tests/conv1d/conv1d.onnx")
        .out_dir("model/named_mpk_gz/")
        .record_type(RecordType::NamedMpkGz)
        .run_from_script();

    ModelGen::new()
        .input("tests/conv1d/conv1d.onnx")
        .out_dir("model/named_mpk_gz_half/")
        .record_type(RecordType::NamedMpkGz)
        .half_precision(true)
        .run_from_script();

    ModelGen::new()
        .input("tests/conv1d/conv1d.onnx")
        .out_dir("model/bincode/")
        .record_type(RecordType::Bincode)
        .run_from_script();

    ModelGen::new()
        .input("tests/conv1d/conv1d.onnx")
        .out_dir("model/bincode_half/")
        .record_type(RecordType::Bincode)
        .half_precision(true)
        .run_from_script();

    ModelGen::new()
        .input("tests/conv1d/conv1d.onnx")
        .out_dir("model/bincode_embedded/")
        .embed_states(true)
        .record_type(RecordType::Bincode)
        .run_from_script();

    ModelGen::new()
        .input("tests/conv1d/conv1d.onnx")
        .out_dir("model/bincode_embedded_half/")
        .embed_states(true)
        .half_precision(true)
        .record_type(RecordType::Bincode)
        .run_from_script();

    // panic!("Purposefully failing build to output logs.");
}
